﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Runtime.InteropServices;
using System.ComponentModel;
using Microsoft.Win32.SafeHandles;
using System.IO;

namespace usbburn
{
    class LowLevelIONt : NtDll, IDisposable
    {
        IntPtr m_handle = IntPtr.Zero;

        public LowLevelIONt(string device, FileAccess fa)
        {
            IO_STATUS_BLOCK st = new IO_STATUS_BLOCK();
            UNICODE_STRING path = new UNICODE_STRING();
            OBJECT_ATTRIBUTES oa = new OBJECT_ATTRIBUTES();

            try
            {
                RtlInitUnicodeString(ref path, device);

                InitializeObjectAttributes(ref oa, path, 0, 0, IntPtr.Zero);

                uint flags = 0;
                switch (fa)
                {
                    case FileAccess.Read:
                        flags = GENERIC_READ;
                        break;
                    case FileAccess.Write:
                        flags = GENERIC_WRITE;
                        break;
                    case FileAccess.ReadWrite:
                        flags = GENERIC_WRITE | GENERIC_READ;
                        break;
                    default:
                        throw new ArgumentException();
                }

                int i = NtOpenFile(ref m_handle, flags | SYNCHRONIZE, ref oa, ref st, 0, FILE_SYNCHRONOUS_IO_NONALERT);

                if (!NT_SUCCESS(i))
                    throw new NtStatusException(i);
            }
            finally
            {
                FreeObjectAttributes(ref oa);
            }
        }

        public void Dispose()
        {
            if (m_handle != IntPtr.Zero)
            {
                NtClose(m_handle);
                m_handle = IntPtr.Zero;
            }
        }

        public byte[] Read(int size)
        {
            byte[] buf = new byte[size];
            int read = 0;

            bool s = ReadFile(m_handle, buf, buf.Length, ref read, IntPtr.Zero);

            if (!s)
                Marshal.ThrowExceptionForHR(Marshal.GetHRForLastWin32Error());

            return buf;
        }

        public void Write(byte[] buf)
        {
            int write = 0;

            bool s = WriteFile(m_handle, buf, buf.Length, ref write, IntPtr.Zero);

            if (!s)
                Marshal.ThrowExceptionForHR(Marshal.GetHRForLastWin32Error());
        }

        ~LowLevelIONt()
        {
            Dispose();
        }
    }
}
